{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Sum of Distances"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Medium"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #array #hash-table #prefix-sum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #数组 #哈希表 #前缀和"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: distance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #等值距离和"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一个下标从 <strong>0</strong> 开始的整数数组 <code>nums</code> 。现有一个长度等于 <code>nums.length</code> 的数组 <code>arr</code> 。对于满足 <code>nums[j] == nums[i]</code> 且 <code>j != i</code> 的所有 <code>j</code> ，<code>arr[i]</code> 等于所有 <code>|i - j|</code> 之和。如果不存在这样的 <code>j</code> ，则令 <code>arr[i]</code> 等于 <code>0</code> 。</p>\n",
    "\n",
    "<p>返回数组<em> </em><code>arr</code><em> 。</em></p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>nums = [1,3,1,1,2]\n",
    "<strong>输出：</strong>[5,0,3,4,0]\n",
    "<strong>解释：</strong>\n",
    "i = 0 ，nums[0] == nums[2] 且 nums[0] == nums[3] 。因此，arr[0] = |0 - 2| + |0 - 3| = 5 。 \n",
    "i = 1 ，arr[1] = 0 因为不存在值等于 3 的其他下标。\n",
    "i = 2 ，nums[2] == nums[0] 且 nums[2] == nums[3] 。因此，arr[2] = |2 - 0| + |2 - 3| = 3 。\n",
    "i = 3 ，nums[3] == nums[0] 且 nums[3] == nums[2] 。因此，arr[3] = |3 - 0| + |3 - 2| = 4 。 \n",
    "i = 4 ，arr[4] = 0 因为不存在值等于 2 的其他下标。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>nums = [0,5,3]\n",
    "<strong>输出：</strong>[0,0,0]\n",
    "<strong>解释：</strong>因为 nums 中的元素互不相同，对于所有 i ，都有 arr[i] = 0 。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= nums.length &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>0 &lt;= nums[i] &lt;= 10<sup>9</sup></code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [sum-of-distances](https://leetcode.cn/problems/sum-of-distances/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [sum-of-distances](https://leetcode.cn/problems/sum-of-distances/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[1,3,1,1,2]', '[0,5,3]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def maxDepth(self, root: Optional[TreeNode]) -> int:\n",
    "        if not root:\n",
    "            return 0\n",
    "        queue = collections.deque([root])\n",
    "        distance = {root: 1}\n",
    "        while queue:\n",
    "            node = queue.popleft()\n",
    "            if node.left:\n",
    "                queue.append(node.left)\n",
    "                distance[node.left] = distance[node] + 1\n",
    "            if node.right:\n",
    "                queue.append(node.right)\n",
    "                distance[node.right] = distance[node] + 1\n",
    "        return  distance[max(distance, key=distance.get)]    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxDistance(self, nums1: List[int], nums2: List[int]) -> int:\n",
    "        m = len(nums1)\n",
    "        n = len(nums2)\n",
    "        i = 0\n",
    "        distance = 0\n",
    "        for j in range(n):\n",
    "            while i < m and nums1[i] > nums2[j]:\n",
    "                i += 1\n",
    "            if i < m:\n",
    "                distance = max(distance,j - i)\n",
    "        return distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxDistance(self, nums1: List[int], nums2: List[int]) -> int:\n",
    "        i = min(len(nums2) - 1, len(nums1) - 1)\n",
    "        max_distance = 0\n",
    "        l, r = i - 1, len(nums2)\n",
    "        while i >= 0:\n",
    "            while l + 1 < r:\n",
    "                m = (l + r) // 2\n",
    "                if nums2[m] >= nums1[i]:\n",
    "                    l = m\n",
    "                else:\n",
    "                    r = m\n",
    "            max_distance = max(max_distance, l - i)\n",
    "            i -= 1\n",
    "            l, r = i - 1, l + 1\n",
    "        return max_distance\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxDistance(self, arrays: List[List[int]]) -> int:\n",
    "        distance=0\n",
    "        self.min=float('inf')\n",
    "        self.max=float('-inf')\n",
    "\n",
    "        \n",
    "        ans=0\n",
    "        for array in arrays:\n",
    "            t_min=float('inf')\n",
    "            t_max=float('-inf')\n",
    "            for i in array:\n",
    "                if i <t_min:\n",
    "                    t_min=i\n",
    "                if i>t_max:\n",
    "                    t_max=i\n",
    "            print(self.max,t_min,t_max,self.min)\n",
    "            ans=max(ans,self.max-t_min,t_max-self.min)\n",
    "            print(ans)\n",
    "            self.max=max(self.max,t_max)\n",
    "            self.min=min(self.min,t_min)\n",
    "    \n",
    "        return ans\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxDistance(self, arrays: List[List[int]]) -> int:\n",
    "        max_v = max([max(i) for i in arrays])\n",
    "        max_v_indexs = [i for i,array in enumerate(arrays) if max_v in array]\n",
    "        min_v = min([min(i) for i in arrays])\n",
    "        min_v_indexs = [i for i,array in enumerate(arrays) if min_v in array]\n",
    "        distance = 0\n",
    "        for i, array in enumerate(arrays):\n",
    "            if max_v not in array or (max_v in array and len(max_v_indexs)>1):\n",
    "                distance = max(abs(max_v - array[0]), distance)\n",
    "            if min_v not in array or (min_v in array and len(min_v_indexs)>1):\n",
    "                distance = max(abs(array[-1]-min_v), distance)\n",
    "        return distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maxDistance(self, arrays: List[List[int]]) -> int:\n",
    "        max_v = max([max(i) for i in arrays])\n",
    "        max_v_indexs = [i for i,array in enumerate(arrays) if max_v in array]\n",
    "        min_v = min([min(i) for i in arrays])\n",
    "        min_v_indexs = [i for i,array in enumerate(arrays) if min_v in array]\n",
    "        distance = 0\n",
    "        for i, array in enumerate(arrays):\n",
    "            if max_v not in array or (max_v in array and len(max_v_indexs)>1):\n",
    "                distance = max(abs(max_v - array[0]), distance)\n",
    "            if min_v not in array or (min_v in array and len(min_v_indexs)>1):\n",
    "                distance = max(abs(array[-1]-min_v), distance)\n",
    "        return distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def maxDistance(self, arrays: List[List[int]]) -> int:\n",
    "        def maxdistance(current_max,current_min,current_result,new_array):\n",
    "            result_candidate_0 = new_array[-1] - current_min \n",
    "            result_candidate_1 = current_max - new_array[0] \n",
    "            current_max = max(current_max,new_array[-1])\n",
    "            current_min = min(current_min,new_array[0])\n",
    "            current_result = max(max(current_result,result_candidate_0),result_candidate_1)\n",
    "            return current_max,current_min,current_result\n",
    "        non_empty_index = [] \n",
    "        for i , array in enumerate(arrays):\n",
    "            if len(array) > 0:\n",
    "                non_empty_index.append(i)\n",
    "        non_empty_arrays = [arrays[i] for i in non_empty_index] \n",
    "        current_max = max(non_empty_arrays[0][-1],non_empty_arrays[1][-1])\n",
    "        current_min = min(non_empty_arrays[0][0],non_empty_arrays[1][0])\n",
    "        current_result = max(non_empty_arrays[0][-1] - non_empty_arrays[1][0], \n",
    "                            non_empty_arrays[1][-1] - non_empty_arrays[0][0] )\n",
    "        for array in non_empty_arrays[2:]:\n",
    "            current_max,current_min,current_result = maxdistance(current_max,current_min,current_result,array)\n",
    "        return current_result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def maximumInvitations(self, favorite: List[int]) -> int:\n",
    "        seen = [False]  * len(favorite)\n",
    "        ans = 0\n",
    "        for i, f in enumerate(favorite):\n",
    "            if not seen[i]:\n",
    "                length = 1\n",
    "                seen[i] = True\n",
    "                distance = {i: 0}\n",
    "                while not seen[f]:\n",
    "                    distance[f] = length\n",
    "                    length += 1\n",
    "                    seen[f] = True\n",
    "                    f = favorite[f]\n",
    "                if f in distance:\n",
    "                    ans = max(ans, length - distance[f])\n",
    "        n = len(favorite)\n",
    "        degree = [0] * n\n",
    "        for i, f in enumerate(favorite):\n",
    "            degree[f] += 1\n",
    "        nodes = set()\n",
    "        for i in range(n):\n",
    "            if i not in nodes and favorite[favorite[i]] == i:\n",
    "                nodes.add(i)\n",
    "                nodes.add(favorite[i])\n",
    "        distance = [1] * n\n",
    "        dq = deque()\n",
    "        for i, d in enumerate(degree):\n",
    "            if d == 0:\n",
    "                dq.append(i)\n",
    "        while dq:\n",
    "            u = dq.popleft()\n",
    "            v = favorite[u]\n",
    "            distance[v] = max(distance[v], distance[u] + 1)\n",
    "            degree[v] -= 1\n",
    "            if degree[v] == 0 and v not in nodes:\n",
    "                dq.append(v)\n",
    "        return max(ans, sum(distance[i] for i in nodes))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def captureForts(self, forts: List[int]) -> int:\n",
    "        distance = []\n",
    "        start = 0\n",
    "        end = 0\n",
    "        target = 1\n",
    "        for i in range(len(forts)):\n",
    "            if target:\n",
    "                if forts[i]==-1 or forts[i]==1:\n",
    "                    start=i\n",
    "                    end=i\n",
    "                    target=0\n",
    "            else:\n",
    "                if forts[i]==-1 or forts[i]==1:\n",
    "                    end = i\n",
    "                    if forts[start]+forts[end]==0:\n",
    "                        distance.append(end-start-1)\n",
    "                    start=i\n",
    "                else:\n",
    "                    end = i\n",
    "\n",
    "        return max(distance) if len(distance)>0 else 0\n",
    "            \n",
    "            \n",
    "                \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def pathsWithMaxScore(self, board: List[str]) -> List[int]:\n",
    "        mod = 10**9+7\n",
    "        n = len(board)\n",
    "        distance = [[0 for j in range(n)] for i in range(n)]\n",
    "        count = [[0 for j in range(n)] for i in range(n)]\n",
    "        count[-1][-1] = 1\n",
    "        for i in range(n-1,-1,-1):\n",
    "            for j in range(n-1,-1,-1):\n",
    "                if board[i][j] != 'S' and board[i][j] != 'X':\n",
    "                    max_val = 0\n",
    "                    temp_count = 0\n",
    "                    if i < n-1:\n",
    "                        if board[i+1][j] != 'X':\n",
    "                            if distance[i+1][j] > max_val:\n",
    "                                max_val = distance[i+1][j]\n",
    "                                temp_count = count[i+1][j]\n",
    "                            elif distance[i+1][j] == max_val:\n",
    "                                temp_count += count[i+1][j]\n",
    "                    if j < n-1:\n",
    "                        if board[i][j+1] != 'X':\n",
    "                            if distance[i][j+1] > max_val:\n",
    "                                max_val = distance[i][j+1]\n",
    "                                temp_count = count[i][j+1]\n",
    "                            elif distance[i][j+1] == max_val:\n",
    "                                temp_count += count[i][j+1]\n",
    "                    if i < n-1 and j < n-1:\n",
    "                        if board[i+1][j+1] != 'X':\n",
    "                            if distance[i+1][j+1] > max_val:\n",
    "                                max_val = distance[i+1][j+1]\n",
    "                                temp_count = count[i+1][j+1]\n",
    "                            elif distance[i+1][j+1] == max_val:\n",
    "                                temp_count += count[i+1][j+1]\n",
    "                    if board[i][j]=='E':\n",
    "                        distance[i][j] = max_val\n",
    "                    else:\n",
    "                        distance[i][j] = (int(board[i][j])+max_val)%mod\n",
    "                    count[i][j] = (count[i][j]+temp_count)%mod\n",
    "        if count[0][0] == 0:\n",
    "            return [0,0]\n",
    "        else:\n",
    "            return [distance[0][0],count[0][0]]\n",
    "                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countRestrictedPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        from collections import defaultdict\n",
    "        import heapq\n",
    "        edge_length = defaultdict(dict)\n",
    "        for edge in edges:\n",
    "            edge_length[edge[0]][edge[1]] = edge[2]\n",
    "            edge_length[edge[1]][edge[0]] = edge[2]\n",
    "        distance = [sys.maxsize] * (n + 1)\n",
    "        distance[n] = 0\n",
    "        seen = set()\n",
    "        queue = [(0, n)]\n",
    "\n",
    "        while queue:\n",
    "            _, src = heapq.heappop(queue)\n",
    "            if src not in seen:\n",
    "                seen.add(src)\n",
    "                for dest, cost in edge_length.get(src).items():\n",
    "                    distance[dest] = min(distance[dest], distance[src] + cost)\n",
    "                    heapq.heappush(queue, (distance[dest], dest))\n",
    "\n",
    "        dist_vertex = [(v, i) for i, v in enumerate(distance) if i > 0]\n",
    "        dist_vertex.sort()\n",
    "\n",
    "        dp = [0] * (n + 1)\n",
    "        dp[n] = 1\n",
    "        for _, src in dist_vertex:\n",
    "            for dst, _ in edge_length.get(src).items():\n",
    "                if distance[src] < distance[dst]:\n",
    "                    dp[dst] += dp[src]\n",
    "\n",
    "        return dp[1] % (10 ** 9 + 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countRestrictedPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        from collections import defaultdict\n",
    "        import heapq\n",
    "        edge_length = defaultdict(dict)\n",
    "        for edge in edges:\n",
    "            edge_length[edge[0]][edge[1]] = edge[2]\n",
    "            edge_length[edge[1]][edge[0]] = edge[2]\n",
    "        distance = [sys.maxsize] * (n + 1)\n",
    "        distance[n] = 0\n",
    "        seen = set()\n",
    "        queue = [(0, n)]\n",
    "\n",
    "        while queue:\n",
    "            _, src = heapq.heappop(queue)\n",
    "            if src not in seen:\n",
    "                seen.add(src)\n",
    "                for dest, cost in edge_length.get(src).items():\n",
    "                    distance[dest] = min(distance[dest], distance[src] + cost)\n",
    "                    heapq.heappush(queue, (distance[dest], dest))\n",
    "\n",
    "        dist_vertex = [(v, i) for i, v in enumerate(distance) if i > 0]\n",
    "        dist_vertex.sort()\n",
    "\n",
    "        dp = [0] * (n + 1)\n",
    "        dp[n] = 1\n",
    "        for _, src in dist_vertex:\n",
    "            for dst, _ in edge_length.get(src).items():\n",
    "                if distance[src] < distance[dst]:\n",
    "                    dp[dst] += dp[src]\n",
    "\n",
    "        return dp[1] % (10 ** 9 + 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countRestrictedPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        def relax(u: int, v: int, w: int) -> None:\n",
    "            if dist[v] > (d := dist[u] + w):\n",
    "                dist[v] = d \n",
    "                heapq.heappush(pq, (d, v))\n",
    "\n",
    "        # 1. O(V+E) build adjacency list of undirected graph\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v, w in edges:\n",
    "            g[u-1].append((v-1, w))\n",
    "            g[v-1].append((u-1, w))\n",
    "\n",
    "        # 2. O((V+E)logV) run heap-optimized Dijkastra's algorithm\n",
    "        start_node = n-1 \n",
    "        dist = [math.inf] * n  \n",
    "        dist[start_node] = 0 \n",
    "        pq = [(0, start_node)]  # min-heap\n",
    "        while pq:\n",
    "            d, u = heapq.heappop(pq)\n",
    "            if d > dist[u]:    \n",
    "                continue\n",
    "       \n",
    "            for v, w in g[u]:\n",
    "                relax(u, v, w)\n",
    "\n",
    "        # 3. DP\n",
    "\n",
    "        # sort vertices by shortest distance ascendingly\n",
    "        vertices = sorted(list(range(n)), key=lambda x: dist[x])\n",
    "\n",
    "        dp = [0] * n \n",
    "        dp[n-1] = 1\n",
    "        for i in vertices:\n",
    "            for j, _ in g[i]:\n",
    "                if dist[i] > dist[j]:\n",
    "                    dp[i] += dp[j]\n",
    "\n",
    "        return dp[0] % (10 ** 9 + 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countRestrictedPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        def Dijkstra(start_node: int, n: int, g: List[List[int]]) -> List[int]:\n",
    "            def relax(u: int, v: int, w: int) -> None:\n",
    "                if dist[v] > (d := dist[u] + w):\n",
    "                    dist[v] = d \n",
    "                    heapq.heappush(pq, (d, v))\n",
    "            \n",
    "            dist = [math.inf] * n  \n",
    "            dist[start_node] = 0 \n",
    "            pq = [(0, start_node)]  # min-heap\n",
    "            while pq:\n",
    "                d, u = heapq.heappop(pq)\n",
    "                if d > dist[u]:    \n",
    "                    continue\n",
    "        \n",
    "                for v, w in g[u]:\n",
    "                    relax(u, v, w)\n",
    "            return dist\n",
    "\n",
    "        # 1. O(V+E) build adjacency list of undirected graph\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u, v, w in edges:\n",
    "            g[u-1].append((v-1, w))\n",
    "            g[v-1].append((u-1, w))\n",
    "\n",
    "        # 2. O((V+E)logV) run heap-optimized Dijkstra's algorithm\n",
    "        dist = Dijkstra(n-1, n, g)\n",
    "\n",
    "        # 3. O(V+E) DP\n",
    "        dp = [0] * n \n",
    "        dp[n-1] = 1\n",
    "\n",
    "        # O(V) sort vertices by shortest distance ascendingly\n",
    "        vertices = sorted(list(range(n)), key=lambda x: dist[x])\n",
    "        \n",
    "        # start from nearest vertex from vertex n \n",
    "        for i in vertices:\n",
    "            for j, _ in g[i]:\n",
    "                if dist[i] > dist[j]:\n",
    "                    dp[i] += dp[j]\n",
    "\n",
    "        return dp[0] % (10 ** 9 + 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countRestrictedPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        graph = defaultdict(list)\n",
    "        for f, t, w in edges:\n",
    "            graph[f].append((t, w))\n",
    "            graph[t].append((f, w))\n",
    "\n",
    "        distance = [float('inf')] * (n + 1)\n",
    "        distance[n] = 0\n",
    "        que, visited = [[0, n]], set()\n",
    "        while que:\n",
    "            _, src = heapq.heappop(que)\n",
    "            if src not in visited:\n",
    "                visited.add(src)\n",
    "                for nxt, w in graph[src]:\n",
    "                    distance[nxt] = min(distance[nxt], distance[src] + w)\n",
    "                    heapq.heappush(que, [distance[nxt], nxt])\n",
    "\n",
    "        sorted_distance = sorted((w, idx) for idx, w in enumerate(distance[1:], 1))\n",
    "        # print(sorted_distance)\n",
    "\n",
    "        dp = [0] * (n + 1)\n",
    "        dp[n] = 1\n",
    "        for _, src in sorted_distance:\n",
    "            for nxt, _ in graph[src]:\n",
    "                if distance[src] < distance[nxt]:\n",
    "                    dp[nxt] += dp[src]\n",
    "        return dp[1] % (10 ** 9 + 7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "from typing import List\n",
    "\n",
    "class Solution:\n",
    "    def countRestrictedPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        mod = 10**9 + 7\n",
    "        graph = [dict() for _ in range(n+1)]\n",
    "        \n",
    "        # 创建邻接表表示的图\n",
    "        for u, v, w in edges:\n",
    "            graph[u][v] = graph[v][u] = w\n",
    "        \n",
    "        # 使用Dijkstra算法从n计算到其他节点的最短距离\n",
    "        distance = [float('inf')] * (n+1)\n",
    "        distance[n] = 0\n",
    "        visited = set()\n",
    "        heap = [(0, n)]\n",
    "        \n",
    "        while heap:\n",
    "            d, node = heapq.heappop(heap)\n",
    "            if node in visited:\n",
    "                continue\n",
    "            visited.add(node)\n",
    "            for neigh, w in graph[node].items():\n",
    "                if distance[neigh] > d + w:\n",
    "                    distance[neigh] = d + w\n",
    "                    heapq.heappush(heap, (distance[neigh], neigh))\n",
    "        \n",
    "        # 使用DFS找出受限路径数\n",
    "        memo = {}\n",
    "        \n",
    "        def dfs(node):\n",
    "            if node == n:\n",
    "                return 1\n",
    "            if node in memo:\n",
    "                return memo[node]\n",
    "            \n",
    "            ans = 0\n",
    "            for neigh in graph[node]:\n",
    "                if distance[neigh] < distance[node]:\n",
    "                    ans += dfs(neigh)\n",
    "                    ans %= mod\n",
    "            memo[node] = ans\n",
    "            return ans\n",
    "        \n",
    "        return dfs(1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def numberOfSubstrings(self, s: str) -> int:\n",
    "\n",
    "        # 利用滑动窗口寻找至少存在abc一次的最短字符串[slow, fast]\n",
    "        # 以slow为开头，结尾从fast 到 (len(s)-1) 由于[slow, fast]至少包含abc，因此后续的都可以\n",
    "        # 只要slow发生变动，对于新的[slow, fast]，其个数为len(s) - fast \n",
    "        slow, fast = 0, 0\n",
    "        res = 0\n",
    "        d = {}\n",
    "        distance = 0\n",
    "        while fast < len(s):\n",
    "            if s[fast] in d:\n",
    "                d[s[fast]] += 1\n",
    "            else:\n",
    "                d[s[fast]] = 1\n",
    "                distance += 1\n",
    "            while distance >= 3:\n",
    "                d[s[slow]] -= 1\n",
    "                if d[s[slow]] == 0:\n",
    "                    del d[s[slow]]\n",
    "                    distance -= 1\n",
    "                slow += 1\n",
    "                res += len(s) - fast\n",
    "            fast += 1\n",
    "        return res\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import heapq\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, roads: List[List[int]]) -> int:\n",
    "        mod = 10**9 + 7\n",
    "        mindis = [10E11]*n\n",
    "        dp = [1]+[0]*(n-1)\n",
    "        pqueue = []\n",
    "        roadgraph = [0 for _ in range(n)]\n",
    "        #构造有向图\n",
    "        for i in range(n):\n",
    "            roadgraph[i] = collections.defaultdict(int)\n",
    "        for u,v,t in roads:\n",
    "            roadgraph[u][v],roadgraph[v][u] = t,t\n",
    "        heapq.heappush(pqueue,[0,0])\n",
    "\n",
    "        #广度优先遍历有向图\n",
    "        while pqueue:\n",
    "            temp = heapq.heappop(pqueue)\n",
    "            curdis = temp[0]\n",
    "            prenode = temp[1]\n",
    "            for node,distance in roadgraph[prenode].items():\n",
    "                if (distance+curdis) < mindis[node]:\n",
    "                    heapq.heappush(pqueue,[distance+curdis,node])\n",
    "                    mindis[node] = distance+curdis\n",
    "                    dp[node] = dp[prenode]\n",
    "                elif (distance+curdis) == mindis[node]:\n",
    "                    dp[node] += dp[prenode]\n",
    "        return dp[n-1]%mod\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    # def dijkstra(self,n,roads):\n",
    "    #     graph=defaultdict(list)\n",
    "    #     for link in roads:\n",
    "    #         x,y,d=roads\n",
    "    #         graph[x].append((y,d))\n",
    "    #         graph[y].append((x,d))\n",
    "    #     que=[(0,0)]#(dist,source)\n",
    "    #     dist=[float('inf')]*n\n",
    "    #     while que:\n",
    "    #         d,cur=heapq.heappop(que)\n",
    "    #         if d>dist[cur]:\n",
    "    #             continue\n",
    "    #         for point,d1 in graph[cur]:\n",
    "    #             if d1>0 and dist[cur]+d1<dist[point]:\n",
    "    #                 dist[point]=dist[cur]+d1\n",
    "    #                 heapq.heappush(que,(dist[point],point))\n",
    "    #     return dist\n",
    "\n",
    "    def countPaths(self, n: int, roads: List[List[int]]) -> int:\n",
    "        MOD = 10**9 + 7\n",
    "        graph = [[] for _ in range(n)]\n",
    "        for u, v, time in roads:\n",
    "            graph[u].append((v, time))\n",
    "            graph[v].append((u, time))\n",
    "        \n",
    "        distances = [float('inf')] * n\n",
    "        path_counts = [0] * n\n",
    "        \n",
    "        distances[0] = 0\n",
    "        path_counts[0] = 1\n",
    "        \n",
    "        heap = [(0, 0)]  # 优先队列，存储距离和节点元组\n",
    "        \n",
    "        while heap:\n",
    "            dist, node = heapq.heappop(heap)\n",
    "            \n",
    "            if dist > distances[node]:\n",
    "                continue\n",
    "            \n",
    "            for neighbor, time in graph[node]:\n",
    "                if dist + time < distances[neighbor]:\n",
    "                    distances[neighbor] = dist + time\n",
    "                    path_counts[neighbor] = path_counts[node]\n",
    "                    heapq.heappush(heap, (dist + time, neighbor))\n",
    "                elif dist + time == distances[neighbor]:\n",
    "                    path_counts[neighbor] = (path_counts[neighbor] + path_counts[node]) % MOD\n",
    "        \n",
    "        return path_counts[n - 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, roads: List[List[int]]) -> int:\n",
    "        graph = [[] for i in range(n)]\n",
    "        for edge in roads:\n",
    "            graph[edge[0]].append((edge[1], edge[2]))\n",
    "            graph[edge[1]].append((edge[0], edge[2]))\n",
    "        def dijkstra(startnode, n):\n",
    "            q = []\n",
    "            distances = [inf for i in range(n)]\n",
    "            distances[startnode] = 0\n",
    "            heapq.heappush(q, (0, startnode))\n",
    "            while len(q) != 0:\n",
    "                temp = heapq.heappop(q)\n",
    "                if temp[0] > distances[temp[1]]:\n",
    "                    continue\n",
    "                for nextnode in graph[temp[1]]:\n",
    "                    distance = temp[0] + nextnode[1]\n",
    "                    if distance < distances[nextnode[0]]:\n",
    "                        distances[nextnode[0]] = distance\n",
    "                        heapq.heappush(q, (distance, nextnode[0]))\n",
    "            return distances\n",
    "        distances = dijkstra(0, n)\n",
    "        print(distances)\n",
    "        g = defaultdict(list)\n",
    "        for x, y, z in roads:\n",
    "            if distances[y] - distances[x] == z:\n",
    "                g[x].append(y)\n",
    "            elif distances[x] - distances[y] == z:\n",
    "                g[y].append(x)\n",
    "        @cache\n",
    "        def dfs(u):\n",
    "            if u == n - 1:\n",
    "                return 1\n",
    "            ret = 0\n",
    "            for v in g[u]:\n",
    "                ret += dfs(v)\n",
    "            return ret % (10**9 + 7)\n",
    "        ans = dfs(0)\n",
    "        dfs.cache_clear()\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, roads: List[List[int]]) -> int:\n",
    "        from collections import defaultdict\n",
    "        graph=defaultdict(list)\n",
    "        for e in roads:\n",
    "            graph[e[0]].append([e[1],e[2]])\n",
    "            graph[e[1]].append([e[0],e[2]])\n",
    "\n",
    "        visited=set()\n",
    "        distance=[[float('inf'),0] if i!=0 else [0,1] for i in range(n)]  #distance[i][0]表示0点到i点的最短距离，distance[i][1]表示0点到i点最短距离的方案数\n",
    "\n",
    "        from heapq import heappush,heappop\n",
    "        heap=[[0,0,0]]  #heap[i][0]表示0点到heap[i][1]点的当前距离，heap[i][1]是目的地节点，heap[i][2]表示目的地节点的父节点\n",
    "        while heap:\n",
    "            node=heappop(heap)\n",
    "            if node[1] not in visited:\n",
    "                visited.add(node[1])\n",
    "                distance[node[1]][0]=node[0]\n",
    "                distance[node[1]][1]=distance[node[2]][1]\n",
    "                for m in graph[node[1]]: \n",
    "                    if m[0] not in visited: heappush(heap,[node[0]+m[1],m[0],node[1]])\n",
    "            else:\n",
    "                if distance[node[1]][0]==node[0]: distance[node[1]][1]=(distance[node[1]][1]+distance[node[2]][1])%1000000007\n",
    "\n",
    "        return distance[-1][1]\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, roads: List[List[int]]) -> int:\n",
    "        # dijkstra算法计算起点到各个点的最短路径\n",
    "        edges = collections.defaultdict(list)\n",
    "        min_distances = [float('inf')]*n\n",
    "        min_distances[0] = 0\n",
    "        for x, y, c in roads:\n",
    "            edges[x].append((y, c))\n",
    "            edges[y].append((x, c))\n",
    "        pq = []\n",
    "        seen = set()\n",
    "        heapq.heappush(pq, (0, 0))\n",
    "        while len(seen)!=n:\n",
    "            min_dis, curr_node = heapq.heappop(pq)\n",
    "            print(curr_node, min_dis)\n",
    "            if curr_node in seen:\n",
    "                continue\n",
    "            else:\n",
    "                seen.add(curr_node)\n",
    "                min_distances[curr_node] = min_dis\n",
    "                for nxt_node, c in edges[curr_node]:\n",
    "                    heapq.heappush(pq, (c+min_dis,nxt_node))\n",
    "        # 构造有向图\n",
    "        graph = collections.defaultdict(list)\n",
    "        for x, y, c in roads:\n",
    "            if min_distances[x]-min_distances[y] == c:\n",
    "                graph[y].append(x)\n",
    "            if min_distances[y]-min_distances[x] == c:\n",
    "                graph[x].append(y)\n",
    "        MOD = 10**9+7\n",
    "        # print(graph)\n",
    "        @cache\n",
    "        def dfs(node):\n",
    "            if node == n-1:\n",
    "                return 1\n",
    "            res = 0\n",
    "            for nxt_node in graph[node]:\n",
    "                res += dfs(nxt_node)\n",
    "            return res%MOD\n",
    "        # print(min_distances)\n",
    "        return dfs(0)\n",
    "\n",
    "\n",
    "# class Solution:\n",
    "#     def countPaths(self, n: int, roads: List[List[int]]) -> int:\n",
    "#         mod = 10**9 + 7\n",
    "        \n",
    "#         dist = [[float(\"inf\")] * n for _ in range(n)]\n",
    "#         for i in range(n):\n",
    "#             dist[i][i] = 0\n",
    "        \n",
    "#         for x, y, z in roads:\n",
    "#             dist[x][y] = dist[y][x] = z\n",
    "        \n",
    "#         # Floyd 算法求解最短路\n",
    "#         # 完成后，dist[0][i] 即为正文部分的 dist[i]\n",
    "#         # for k in range(n):\n",
    "#         #     for i in range(n):\n",
    "#         #         for j in range(n):\n",
    "#         #             dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])\n",
    "\n",
    "#         # Dijkstra 算法求解最短路\n",
    "#         # 完成后，dist[0][i] 即为正文部分的 dist[i]\n",
    "#         seen = set()\n",
    "#         for _ in range(n - 1):\n",
    "#             u = None\n",
    "#             for i in range(n):\n",
    "#                 if i not in seen and (not u or dist[0][i] < dist[0][u]):\n",
    "#                     u = i\n",
    "#             seen.add(u)\n",
    "#             for i in range(n):\n",
    "#                 dist[0][i] = min(dist[0][i], dist[0][u] + dist[u][i])\n",
    "\n",
    "#         # 构造图 G\n",
    "#         g = defaultdict(list)\n",
    "#         for x, y, z in roads:\n",
    "#             if dist[0][y] - dist[0][x] == z:\n",
    "#                 g[x].append(y)\n",
    "#             elif dist[0][x] - dist[0][y] == z:\n",
    "#                 g[y].append(x)\n",
    "\n",
    "#         @cache\n",
    "#         def dfs(u: int) -> int:\n",
    "#             if u == n - 1:\n",
    "#                 return 1\n",
    "\n",
    "#             ret = 0\n",
    "#             for v in g[u]:\n",
    "#                 ret += dfs(v)\n",
    "#             return ret % mod\n",
    "        \n",
    "#         ans = dfs(0)\n",
    "#         dfs.cache_clear()\n",
    "#         print(dist[0])\n",
    "#         return ans\n",
    "        \n",
    "            \n",
    "\n",
    "        \n",
    "\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, roads: List[List[int]]) -> int:\n",
    "        MOD = 10 ** 9 + 7\n",
    "        graph = collections.defaultdict(set)\n",
    "        for u, v, t in roads:\n",
    "            graph[u].add((v, t))\n",
    "            graph[v].add((u, t))\n",
    "        distance = [-1] * n\n",
    "        pq = [(0,0)]\n",
    "\n",
    "        while pq:\n",
    "            d, node = heapq.heappop(pq)\n",
    "            if distance[node] >= 0:\n",
    "                continue\n",
    "            distance[node] = d\n",
    "            for nxt, t in graph[node]:\n",
    "                if distance[nxt] >= 0:\n",
    "                    continue\n",
    "                heapq.heappush(pq, (t + d, nxt))\n",
    "        # print(distance)\n",
    "\n",
    "        memo = {}\n",
    "        def dfs(node, d):\n",
    "            if distance[node] != d:\n",
    "                return 0\n",
    "            if node == 0:\n",
    "                return 1\n",
    "            \n",
    "            if node in memo:\n",
    "                return memo[node]\n",
    "            val = 0\n",
    "            for nxt, t in graph[node]:\n",
    "                val += dfs(nxt, d - t)\n",
    "                val %= MOD\n",
    "            memo[node] = val\n",
    "            return val\n",
    "        return dfs(n-1, distance[-1]) % MOD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        parent_children = {}\n",
    "        for e in edges:\n",
    "            children = parent_children.get(e[0], [])\n",
    "            children.append(e[1])\n",
    "            parent_children[e[0]] = children\n",
    "\n",
    "            children = parent_children.get(e[1], [])\n",
    "            children.append(e[0])\n",
    "            parent_children[e[1]] = children\n",
    "\n",
    "        ready_count = [0 for i in range(n)]\n",
    "        distance = {}\n",
    "\n",
    "        queue = []\n",
    "        for parent in parent_children:\n",
    "            if len(parent_children[parent]) == 1:\n",
    "                queue.append(parent)\n",
    "\n",
    "        res = [0 for i in range(n)]\n",
    "\n",
    "        start = 0\n",
    "        while start < len(queue):\n",
    "            node_id = queue[start]\n",
    "            start += 1\n",
    "            children = parent_children[node_id]\n",
    "            sum_dist = 0\n",
    "            sum_child = 0\n",
    "            miss_child = None\n",
    "            for child in children:\n",
    "                if child * n + node_id in distance:\n",
    "                    [dist, n_child] = distance.get(child * n + node_id)\n",
    "                    sum_dist += dist\n",
    "                    sum_child += n_child\n",
    "                else:\n",
    "                    miss_child = child\n",
    "\n",
    "            if ready_count[node_id] == len(children) - 1:\n",
    "                sum_dist += sum_child + 1\n",
    "                distance[node_id * n + miss_child] = [sum_dist, sum_child + 1]\n",
    "                ready_count[miss_child] += 1\n",
    "                if ready_count[miss_child] >= len(parent_children[miss_child]) - 1:\n",
    "                    queue.append(miss_child)\n",
    "            else:\n",
    "                #self.show(n, distance)\n",
    "                res[node_id] = sum_dist\n",
    "                for child in children:\n",
    "                    if res[child] == 0:\n",
    "                        [dist, n_child] = distance[child * n + node_id]\n",
    "                        distance[node_id * n + child] = [sum_dist - dist + sum_child - n_child + 1, sum_child - n_child + 1]\n",
    "                        ready_count[child] += 1\n",
    "                        if ready_count[child] >= len(parent_children[child]) - 1:\n",
    "                            queue.append(child)\n",
    "\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def dfs_forward(self, point, root, forwards, links):\n",
    "        # forwards[point] = [1, 0]\n",
    "        forwards[point] = 1\n",
    "        distance = 0\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                num, dis = self.dfs_forward(p, point, forwards, links)\n",
    "                # forwards[point][0] += forwards[p][0] # num\n",
    "                # forwards[point][1] += (forwards[p][0] + forwards[p][1]) # distance_sum\n",
    "                forwards[point] += num\n",
    "                distance += num + dis\n",
    "        return forwards[point], distance\n",
    "\n",
    "    def dfs_backward(self, point, root, backwards, forwards, links):\n",
    "        # + (n-k) - k \n",
    "        backwards[point] = backwards[root] + len(backwards) - 2 * forwards[point]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_backward(p, point, backwards, forwards, links)\n",
    "        return  \n",
    "\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if len(edges) == 0:\n",
    "            return [0]\n",
    "        # links = collections.defaultdict(list)\n",
    "        links = [[] for _ in range(n)]\n",
    "        for p1, p2 in edges:\n",
    "            links[p1].append(p2)\n",
    "            links[p2].append(p1)\n",
    "        p0 = edges[0][0]\n",
    "        forwards = [0 for i in range(n)]\n",
    "        backwards = [0 for i in range(n)]\n",
    "        n_total, backwards[p0] = self.dfs_forward(p0, -1, forwards, links)\n",
    "\n",
    "        for p in links[p0]:\n",
    "            self.dfs_backward(p, p0, backwards, forwards, links)\n",
    "        return backwards\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def dfs_forward(self, point, root, forwards, links):\n",
    "        forwards[point] = [1, 0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_forward(p, point, forwards, links)\n",
    "                forwards[point][0] += forwards[p][0] # num\n",
    "                forwards[point][1] += (forwards[p][0] + forwards[p][1]) # distance_sum\n",
    "        return \n",
    "    def dfs_backward(self, point, root, backwards, forwards, links):\n",
    "        # + (n-k) - k \n",
    "        backwards[point] = backwards[root] + len(backwards) - 2 * forwards[point][0]\n",
    "        for p in links[point]:\n",
    "            if p != root:\n",
    "                self.dfs_backward(p, point, backwards, forwards, links)\n",
    "        return  \n",
    "\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        if len(edges) == 0:\n",
    "            return [0]\n",
    "        # links = collections.defaultdict(list)\n",
    "        links = [[] for _ in range(n)]\n",
    "        for p1, p2 in edges:\n",
    "            links[p1].append(p2)\n",
    "            links[p2].append(p1)\n",
    "        p0 = edges[0][0]\n",
    "        forwards = [0 for i in range(n)]\n",
    "        self.dfs_forward(p0, -1, forwards, links)\n",
    "        backwards = [0 for i in range(n)]\n",
    "        backwards[p0] = forwards[p0][1]\n",
    "        for p in links[p0]:\n",
    "            self.dfs_backward(p, p0, backwards, forwards, links)\n",
    "        return backwards\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "\n",
    "\n",
    "        tree = [[] for _ in range(n)]\n",
    "\n",
    "        for edge in edges:\n",
    "            a, b = edge[0], edge[1]\n",
    "            tree[a].append(b) \n",
    "            tree[b].append(a) \n",
    "\n",
    "    \n",
    "        # 以0为root 求出每个子树的size\n",
    "        subtree = [1]*n\n",
    "        parents = [-1]*n\n",
    "        def dfs(root, parent):\n",
    "            for child in tree[root]:\n",
    "                if child == parent:\n",
    "                    continue\n",
    "                parents[child] = root\n",
    "                subtree[root] += dfs(child, root)\n",
    "            return subtree[root]\n",
    "        dfs(0, -1)\n",
    "        # print(tree)\n",
    "        # print()\n",
    "        # print(subtree)\n",
    "        # print(parents)\n",
    "\n",
    "        distance = [0]*n\n",
    "        def dfs2(root, path, parent):\n",
    "            res = 0\n",
    "            for i in tree[root]:\n",
    "                if i == parent:\n",
    "                    continue\n",
    "                distance[i] = path+1\n",
    "                res += dfs2(i, path+1, root)\n",
    "            return res\n",
    "        dfs2(0, 0, -1)\n",
    "\n",
    "        res = [0]*(n+1)\n",
    "        res[0] = sum(distance)\n",
    "\n",
    "        # 应该顺着parent的顺序更新\n",
    "        # for i in range(1, n):\n",
    "        #     s = subtree[i]\n",
    "        #     res[i] = res[parents[i]] - s + (n-s)\n",
    "\n",
    "        def dfs3(root, p):\n",
    "            # print(root)\n",
    "            for child in tree[root]:\n",
    "                if child == p:\n",
    "                    continue\n",
    "                s = subtree[child]\n",
    "                \n",
    "                res[child] = res[root] - s  + (n-s)\n",
    "                # print(\"-\", child, res[root], res[child])\n",
    "                dfs3(child, root)\n",
    "        # for c in tree[0]:\n",
    "        #     dfs3(c, 0)\n",
    "        dfs3(0, -1)\n",
    "        return res[:n]\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        self.n = n\n",
    "        self.tree_size = {i:0 for i in range(n)}\n",
    "        self.graph = {i:set([]) for i in range(n)}\n",
    "        for edge in edges:\n",
    "            x,y = edge\n",
    "            self.graph[x].add(y)\n",
    "            self.graph[y].add(x)\n",
    "        self.ans = [0]*n\n",
    "        self.dfs(0, set([0]), 1)\n",
    "        # tree dp\n",
    "        self.tree_dp(0, set([0]))\n",
    "        return self.ans\n",
    "\n",
    "\n",
    "    \n",
    "    def dfs(self, root, visited, level):\n",
    "        ans = 1\n",
    "        for next_node in self.graph[root]:\n",
    "            if next_node not in visited:\n",
    "                visited.add(next_node)\n",
    "                self.ans[0] += level\n",
    "                ans += self.dfs(next_node, visited, level+1)\n",
    "        self.tree_size[root] = ans\n",
    "        return ans\n",
    "    \n",
    "    def tree_dp(self, root, visited):\n",
    "        for next_node in self.graph[root]:\n",
    "            if next_node not in visited:\n",
    "                # sum_subtree = sum dist(root, node_j) for every subtree node j based on the root\n",
    "                # sum_other = sum dist(root, node_i) for every node i exclude from the subtree base on the root\n",
    "                # sum_distance = sum_subtree + sum_other\n",
    "                # let x be the parent of y\n",
    "                # sum_distance(y) = sum_subtree_y + sum_other_y\n",
    "                # sum_subtree_y = sum_subtree_x - size_y\n",
    "                # sum_other_y = sum_other_x + n - size_y(all excluded nodes dist + 1)\n",
    "                self.ans[next_node] = self.ans[root] + self.n - 2*self.tree_size[next_node]\n",
    "                visited.add(next_node)\n",
    "                self.tree_dp(next_node, visited)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:\n",
    "        def f(node):\n",
    "            for child in children[node]:\n",
    "                f(child)\n",
    "\n",
    "        down_counts = [1] * n  # 当前节点下方节点的个数(包括自己)\n",
    "        up_counts = [0] * n  # 当前节点上方节点的个数(包括父节点的其他子树节点，不包括自己)\n",
    "        depths = [0] * n  # 深度。0节点深度为0\n",
    "        distances = [0] * n  # 距离和\n",
    "\n",
    "        children = defaultdict(set)\n",
    "        for f, t in edges:\n",
    "            children[f].add(t)\n",
    "            children[t].add(f)\n",
    "\n",
    "        def clear_fathers(node):\n",
    "            for child in children[node]:\n",
    "                children[child].discard(node)\n",
    "                clear_fathers(child)\n",
    "\n",
    "        def cal_down_counts(node):\n",
    "            cnt = 1\n",
    "            for child in children[node]:\n",
    "                cnt += cal_down_counts(child)\n",
    "            down_counts[node] = cnt\n",
    "            return cnt\n",
    "\n",
    "        def cal_up_counts(node):\n",
    "            up_count = up_counts[node]\n",
    "            down_count = down_counts[node]\n",
    "            for child in children[node]:\n",
    "                new_up_node_cnt = down_count - down_counts[child]\n",
    "                up_counts[child] = up_count + new_up_node_cnt\n",
    "                cal_up_counts(child)\n",
    "\n",
    "        def cal_depths(node, depth):\n",
    "            depths[node] = depth\n",
    "            for child in children[node]:\n",
    "                cal_depths(child, depth + 1)\n",
    "\n",
    "        def cal_distances(node):\n",
    "            dist = distances[node]\n",
    "            for child in children[node]:\n",
    "                delta = up_counts[child] - down_counts[child]\n",
    "                distances[child] = dist + delta\n",
    "                cal_distances(child)\n",
    "\n",
    "        clear_fathers(0)\n",
    "        cal_down_counts(0)\n",
    "        cal_up_counts(0)\n",
    "        cal_depths(0, 0)\n",
    "        distances[0] = sum(depths)\n",
    "        cal_distances(0)\n",
    "\n",
    "        return distances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        nums=[ [nums[i],i] for i in range(len(nums))]\n",
    "\n",
    "        nums.sort(key=lambda x:(x[0],x[1]))\n",
    "\n",
    "        l=0\n",
    "        r=0\n",
    "\n",
    "        arr=[0 for i in range(len(nums))]\n",
    "\n",
    "        print(nums)\n",
    "        while(r<len(nums)):\n",
    "            if(nums[r][0]==nums[l][0]):\n",
    "                r+=1\n",
    "                if(r==len(nums)):\n",
    "                    sum_=sum([ nums[i][1] for i in range(l,r)])\n",
    "                    pre_sum=0\n",
    "                    for k in range(l,r):\n",
    "                        post_sum=sum_-pre_sum-nums[k][1]\n",
    "                        #print(f\"pre {pre_sum} post sum {post_sum}\")\n",
    "                        arr[nums[k][1]]=(k-l)*nums[k][1]-pre_sum+post_sum-(r-1-k)*nums[k][1]\n",
    "                        pre_sum+=nums[k][1]\n",
    "                    break\n",
    "            elif(r==len(nums) or nums[r][0]!=nums[l][0]):\n",
    "                if(r-l==1):\n",
    "                    arr[nums[l][1]]=0\n",
    "                else:\n",
    "                    sum_=sum([ nums[i][1] for i in range(l,r)])\n",
    "                    pre_sum=0\n",
    "                    for k in range(l,r):\n",
    "                        post_sum=sum_-pre_sum-nums[k][1]\n",
    "                        #print(f\"pre {pre_sum} post sum {post_sum}\")\n",
    "                        arr[nums[k][1]]=(k-l)*nums[k][1]-pre_sum+post_sum-(r-1-k)*nums[k][1]\n",
    "                        pre_sum+=nums[k][1]\n",
    "                        \n",
    "                l=r \n",
    "                \n",
    "        \n",
    "        return arr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        dic = {}\n",
    "        arr = [0] * len(nums)\n",
    "        for i,c in enumerate(nums) :\n",
    "            if c in dic:\n",
    "                dic[c].append(i)\n",
    "            else :\n",
    "                dic[c] = [i]\n",
    "        for k, v in dic.items() :\n",
    "            rt1 = 0 \n",
    "            rt2 = sum(v)\n",
    "            t1 = 0 \n",
    "            t2 = len(v)\n",
    "            for j in v :\n",
    "                arr[j]  = j*t1 - rt1 + rt2 - j*t2 \n",
    "                rt1 += j\n",
    "                rt2 -= j\n",
    "                t1 +=1\n",
    "                t2 -= 1\n",
    "        return arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        arr = [0] * len(nums)\n",
    "        hashdict = {}\n",
    "        for i in range(len(nums)):\n",
    "            num = nums[i]\n",
    "            if num not in hashdict:\n",
    "                hashdict[num] = [i]\n",
    "            else:\n",
    "                hashdict[num].append(i)\n",
    "        for key in hashdict:\n",
    "            keyValues = hashdict[key]\n",
    "            n = len(keyValues)\n",
    "            for i in range(1,n):\n",
    "                arr[keyValues[0]] += abs(keyValues[0] - keyValues[i])\n",
    "            for i in range(1,n):\n",
    "                change = abs(keyValues[i] - keyValues[i - 1])\n",
    "                arr[keyValues[i]] = arr[keyValues[i - 1]] - (n - 2 * i) * change\n",
    "        \n",
    "        return arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        arr = [0] * len(nums)\n",
    "        hashdict = {}\n",
    "        for i in range(len(nums)):\n",
    "            num = nums[i]\n",
    "            if num not in hashdict:\n",
    "                hashdict[num] = [i]\n",
    "            else:\n",
    "                hashdict[num].append(i)\n",
    "        for key in hashdict:\n",
    "            keyValues = hashdict[key]\n",
    "            for i in range(1,len(keyValues)):\n",
    "                arr[keyValues[0]] += abs(keyValues[0] - keyValues[i])\n",
    "            for i in range(1,len(keyValues)):\n",
    "                sum = len(keyValues) - 1\n",
    "                add = sum - i\n",
    "                sub = sum - add - 1\n",
    "                change = abs(keyValues[i] - keyValues[i - 1])\n",
    "                arr[keyValues[i]] = arr[keyValues[i - 1]] - (add - sub) * change\n",
    "        \n",
    "        return arr\n",
    "\n",
    "\n",
    "\n",
    "        \n",
    "        return arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        arr = [0] * len(nums)\n",
    "        hashdict = {}\n",
    "        for i in range(len(nums)):\n",
    "            num = nums[i]\n",
    "            if num not in hashdict:\n",
    "                hashdict[num] = [i]\n",
    "            else:\n",
    "                hashdict[num].append(i)\n",
    "        for key in hashdict:\n",
    "            keyValues = hashdict[key]\n",
    "            n = len(keyValues)\n",
    "            arr[keyValues[0]] = sum(abs(keyValues[0] - keyValues[i]) for i in range(1,n) )\n",
    "            for i in range(1,n):\n",
    "                change = abs(keyValues[i] - keyValues[i - 1])\n",
    "                arr[keyValues[i]] = arr[keyValues[i - 1]] - (n - 2 * i) * change       \n",
    "        return arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n=len(nums)\n",
    "        ad={}\n",
    "        for i,c in enumerate(nums):\n",
    "            if c in ad:\n",
    "                ad[c].append(i)\n",
    "            else:\n",
    "                ad[c]=[i]\n",
    "        ans=[0]*n\n",
    "        for arr in ad.values():\n",
    "            if len(arr)==1:\n",
    "                continue\n",
    "            x=len(arr)\n",
    "            dp=[0]*len(arr)\n",
    "            print(arr)\n",
    "            dp[0]=arr[0]\n",
    "            for j in range(1,x):\n",
    "                dp[j]=arr[j]+dp[j-1]\n",
    "                \n",
    "            for k,target in enumerate(arr):\n",
    "                left=(k+1)*target-dp[k]\n",
    "                right=dp[-1]-dp[k]-target*(x-k-1)\n",
    "                ans[target]=left+right\n",
    "                #print(left,right,k,target)\n",
    "        #l=arr.index(i)\n",
    "           # print(dp)\n",
    "           ##left=(l+1)*i-dp[l]\n",
    "           # \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        arr = [0] * len(nums)\n",
    "        hashdict = {}\n",
    "        for i in range(len(nums)):\n",
    "            num = nums[i]\n",
    "            if num not in hashdict:\n",
    "                hashdict[num] = [i]\n",
    "            else:\n",
    "                hashdict[num].append(i)\n",
    "        for keyValues in hashdict.values():\n",
    "            n = len(keyValues)\n",
    "            arr[keyValues[0]] = sum(abs(keyValues[0] - keyValues[i]) for i in range(1,n) )\n",
    "            for i in range(1,n):\n",
    "                arr[keyValues[i]] = arr[keyValues[i - 1]] - (n - 2 * i) * abs(keyValues[i] - keyValues[i - 1])       \n",
    "        return arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        arr = [0 for _ in nums]\n",
    "        sum = dict()\n",
    "        cnt = dict()\n",
    "        for i in range(0, len(nums)):\n",
    "            if (nums[i] not in sum):\n",
    "                sum[nums[i]] = i\n",
    "                cnt[nums[i]] = 1\n",
    "            else:\n",
    "                arr[i] += i * cnt[nums[i]] - sum[nums[i]]\n",
    "                sum[nums[i]] += i\n",
    "                cnt[nums[i]] += 1\n",
    "\n",
    "        sum = {}\n",
    "        cnt = {}\n",
    "        for i in range(len(nums) - 1, -1, -1):\n",
    "            if (nums[i] not in sum):\n",
    "                sum[nums[i]] = i\n",
    "                cnt[nums[i]] = 1\n",
    "            else:\n",
    "                arr[i] += sum[nums[i]] - i * cnt[nums[i]]\n",
    "                sum[nums[i]] += i\n",
    "                cnt[nums[i]] += 1\n",
    "\n",
    "        return arr\n",
    "             \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        mp = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            mp[x].append(i)\n",
    "        ans = [0] * n\n",
    "        for g in mp.values():\n",
    "            #[0...i - 1] [i... n - 1]\n",
    "            # i             l - i\n",
    "            l = len(g)\n",
    "            s = [0] * (l + 1)\n",
    "            for i in range(l):\n",
    "                s[i + 1] = s[i] + g[i]\n",
    "            for i in range(l):\n",
    "                x = g[i]\n",
    "                ans[x] = x * i - s[i] + s[l] - s[i] - (l - i) * x\n",
    "        return ans\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        arr = [0] * len(nums)\n",
    "        hashdict = {}\n",
    "        for i in range(len(nums)):\n",
    "            num = nums[i]\n",
    "            if num not in hashdict:\n",
    "                hashdict[num] = [i]\n",
    "            else:\n",
    "                hashdict[num].append(i)\n",
    "        for key in hashdict:\n",
    "            keyValues = hashdict[key]\n",
    "            for i in range(1,len(keyValues)):\n",
    "                arr[keyValues[0]] += abs(keyValues[0] - keyValues[i])\n",
    "            for i in range(1,len(keyValues)):\n",
    "                sum = len(keyValues) - 1\n",
    "                add = sum - i\n",
    "                sub = sum - add - 1\n",
    "                change = abs(keyValues[i] - keyValues[i - 1])\n",
    "                arr[keyValues[i]] = arr[keyValues[i - 1]] - (sum - 2 * i + 1) * change\n",
    "        \n",
    "        return arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n=len(nums)\n",
    "        a=[0]*n\n",
    "        c=Counter()\n",
    "        s=Counter()\n",
    "        for i in range(n):\n",
    "            x=nums[i]\n",
    "            a[i]+=c[x]*i-s[x]\n",
    "            s[x]+=i\n",
    "            c[x]+=1\n",
    "        c=Counter()\n",
    "        s=Counter()\n",
    "        for i in range(n-1,-1,-1):\n",
    "            x=nums[i]\n",
    "            a[i]+=s[x]-c[x]*i\n",
    "            s[x]+=i\n",
    "            c[x]+=1        \n",
    "        return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        numDict = collections.defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            numDict[num].append(i)\n",
    "        \n",
    "        for num in numDict:\n",
    "            if len(numDict[num]) == 1:\n",
    "                nums[numDict[num][0]] = 0\n",
    "            else:\n",
    "                post = sum(numDict[num])\n",
    "                pre = 0\n",
    "                n = len(numDict[num])\n",
    "                for i in range(n):\n",
    "                    post -= numDict[num][i]\n",
    "                    nums[numDict[num][i]] = post - pre - (n - 2*i - 1) * numDict[num][i]\n",
    "                    pre += numDict[num][i]\n",
    "        return nums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        numDict = collections.defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            numDict[num].append(i)\n",
    "        \n",
    "        for num in numDict:\n",
    "            if len(numDict[num]) == 1:\n",
    "                nums[numDict[num][0]] = 0\n",
    "            else:\n",
    "                post = sum(numDict[num])\n",
    "                pre = 0\n",
    "                n = len(numDict[num])\n",
    "                for i in range(n):\n",
    "                    post -= numDict[num][i]\n",
    "                    nums[numDict[num][i]] = post - pre - (n - 2*i - 1) * numDict[num][i]\n",
    "                    pre += numDict[num][i]\n",
    "        return nums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n=len(nums)\n",
    "        ad={}\n",
    "        for i,c in enumerate(nums):\n",
    "            if c in ad:\n",
    "                ad[c].append(i)\n",
    "            else:\n",
    "                ad[c]=[i]\n",
    "        ans=[0]*n\n",
    "        for arr in ad.values():\n",
    "            if len(arr)==1:\n",
    "                continue\n",
    "            x=len(arr)\n",
    "            dp=[0]*len(arr)\n",
    "            #print(arr)\n",
    "            dp[0]=arr[0]\n",
    "            for j in range(1,x):\n",
    "                dp[j]=arr[j]+dp[j-1]\n",
    "            for k,target in enumerate(arr):\n",
    "                left=(k+1)*target-dp[k]\n",
    "                right=dp[-1]-dp[k]-target*(x-k-1)\n",
    "                ans[target]=left+right\n",
    "                #print(left,right,k,target)\n",
    "        #l=arr.index(i)\n",
    "           # print(dp)\n",
    "           ##left=(l+1)*i-dp[l]\n",
    "           # \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        mp = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            mp[x].append(i)\n",
    "        ans = [0] * n\n",
    "        s = [0] * (n + 1)\n",
    "        for g in mp.values():\n",
    "            #[0...i - 1] [i... n - 1]\n",
    "            # i             l - i\n",
    "            l = len(g)\n",
    "            for i in range(l):\n",
    "                s[i + 1] = s[i] + g[i]\n",
    "            for i in range(l):\n",
    "                x = g[i]\n",
    "                ans[x] = x * i - s[i] + s[l] - s[i] - (l - i) * x\n",
    "        return ans\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        numDict = collections.defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            numDict[num].append(i)\n",
    "        \n",
    "        for num in numDict:\n",
    "            if len(numDict[num]) == 1:\n",
    "                nums[numDict[num][0]] = 0\n",
    "            else:\n",
    "                post = sum(numDict[num])\n",
    "                pre = 0\n",
    "                n = len(numDict[num])\n",
    "                for i in range(n):\n",
    "                    post -= numDict[num][i]\n",
    "                    nums[numDict[num][i]] = post - pre - (n - 2*i - 1) * numDict[num][i]\n",
    "                    pre += numDict[num][i]\n",
    "        return nums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        numDict = collections.defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            numDict[num].append(i)\n",
    "        \n",
    "        for num in numDict:\n",
    "            if len(numDict[num]) == 1:\n",
    "                nums[numDict[num][0]] = 0\n",
    "            else:\n",
    "                post = sum(numDict[num])\n",
    "                pre = 0\n",
    "                n = len(numDict[num])\n",
    "                for i in range(n):\n",
    "                    post -= numDict[num][i]\n",
    "                    nums[numDict[num][i]] = post - pre - (n - 2*i - 1) * numDict[num][i]\n",
    "                    pre += numDict[num][i]\n",
    "        return nums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        locs = defaultdict(list)\n",
    "        for i,x in enumerate(nums): locs[x].append(i)\n",
    "        ans = [0] * len(nums)\n",
    "        for x,l in locs.items():\n",
    "            if len(l) == 1: continue\n",
    "            s,n,acc = sum(l),len(l),0\n",
    "            for c,i in enumerate(l):\n",
    "                ans[i] = (2*c-n)*i + s - acc*2\n",
    "                acc += i\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        mp=defaultdict(list)\n",
    "        for i,x in enumerate(nums) :\n",
    "            mp[x].append(i)\n",
    "        n=len(nums)\n",
    "        ans=[0]*n\n",
    "        s=[0]*(n+1)\n",
    "        #x是下标数组\n",
    "        for x in mp.values() :\n",
    "            n=len(x)\n",
    "            for i in range(1,n+1) :\n",
    "                s[i]=s[i-1]+x[i-1]\n",
    "            for i,val in enumerate(x) :\n",
    "                ans[val]=i*val-s[i]+s[n]-s[i+1]-(n-1-i)*val\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, a: List[int]) -> List[int]:\n",
    "        n=len(a)\n",
    "        h=defaultdict(list)\n",
    "        for i,j in enumerate(a):\n",
    "            h[j].append(i)\n",
    "        # print(h)\n",
    "        ans=[0]*n\n",
    "        for i in h.values():\n",
    "            m=len(i)\n",
    "            b=list(accumulate(i))\n",
    "            for j,k in enumerate(i):\n",
    "                # print(i,j,k)\n",
    "                ans[k]=(k*(j+1))-b[j]+(b[-1]-b[j])-k*(m-j-1)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        #n=len(nums)\n",
    "        #ans=[0]*n\n",
    "        #group=defaultdict(list)\n",
    "        ##{1:[0,2,3];3:[1];2:[4]}\n",
    "        #for i,x in enumerate(nums):\n",
    "        #    group[x].append(i)\n",
    "        #for a in group.values():\n",
    "        #    s=list(accumulate(a,initial=0))\n",
    "        #    #[0,0,2,5]\n",
    "        #    m=len(a)\n",
    "        #    for i in a:\n",
    "        #        b=bisect_left(a,i)\n",
    "        #        left=i*b-s[b]\n",
    "        #        right=s[m]-s[b]-i*(m-b)\n",
    "        #        ans[i]=left+right\n",
    "        #return ans\n",
    "\n",
    "        n=len(nums)\n",
    "        ans=[0]*n\n",
    "        group=defaultdict(list)\n",
    "        #{1:[0,2,3];3:[1];2:[4]}\n",
    "        for i,x in enumerate(nums):\n",
    "            group[x].append(i)\n",
    "        for a in group.values():\n",
    "            m=len(a)\n",
    "            s=0\n",
    "            first=a[0]\n",
    "            for i in range(1,m):\n",
    "                s += abs(first-a[i])\n",
    "            ans[a[0]]=s\n",
    "            for i in range(1,m):\n",
    "                ans[a[i]] = s+i*(a[i]-a[i-1])-(m-i)*(a[i]-a[i-1])\n",
    "                s += i*(a[i]-a[i-1])-(m-i)*(a[i]-a[i-1])\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        ex = defaultdict(list)\n",
    "        ans = [0]*len(nums)\n",
    "        for idx,v in enumerate(nums):\n",
    "            ex[v].append(idx)\n",
    "        for _,l in ex.items():\n",
    "            cur_sum = sum(l)\n",
    "            pre = 0\n",
    "            for cnt,idx in enumerate(l):\n",
    "                ans[idx]=(cur_sum-idx-pre)-(len(l)-1-cnt)*idx+cnt*idx-pre\n",
    "                pre+=idx\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        l = len(nums)\n",
    "        arr = []\n",
    "        for i in range(l):\n",
    "            arr.append(0)\n",
    "        r = {}\n",
    "        for i in range(l):\n",
    "            if nums[i] not in r.keys():\n",
    "                r[nums[i]] = [i]\n",
    "            else:\n",
    "                r[nums[i]].append(i)\n",
    "        for x in r.keys():\n",
    "            lr = len(r[x])\n",
    "            if lr>1:\n",
    "                for i in range(len(r[x])):\n",
    "                    if i == 0:\n",
    "                        for j in range(i,len(r[x])):\n",
    "                            c = abs(r[x][i] - r[x][j])\n",
    "                            arr[r[x][i]] += c\n",
    "                            # arr[r[x][j]] += c\n",
    "                    else:\n",
    "                        arr[r[x][i]] = arr[r[x][i-1]]-(lr-i)*(r[x][i]-r[x][i-1])+(i)*(r[x][i]-r[x][i-1])\n",
    "        return arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        arr = [0] * n\n",
    "        cnt = {}\n",
    "        for i, x in enumerate(nums):\n",
    "            if x not in cnt:\n",
    "                cnt[x] = [i]\n",
    "            else:\n",
    "                cnt[x].append(i)\n",
    "        \n",
    "        for pos in cnt.values():\n",
    "            l = len(pos)\n",
    "            s = list(accumulate(pos, initial = 0))\n",
    "            for j, target in enumerate(pos):\n",
    "                left = target * j - s[j]\n",
    "                right = s[-1] - s[j] - target * (l - j)\n",
    "                arr[target] = left + right\n",
    "        return arr\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "      def diff(nums):\n",
    "        n = len(nums)\n",
    "        prefix = list(accumulate(nums))\n",
    "\n",
    "        res = [0] * n\n",
    "        res[0] = (n - 1) * (-nums[0]) + prefix[-1] - prefix[0]\n",
    "\n",
    "        for i in range(1, n):\n",
    "          res[i] = prefix[-1] - prefix[i] - prefix[i - 1] + (i * 2 - n + 1) * nums[i]\n",
    "\n",
    "        return res\n",
    "\n",
    "      indices = defaultdict(list)\n",
    "\n",
    "      for i, num in enumerate(nums):\n",
    "        indices[num].append(i)\n",
    "\n",
    "      # print(indices)\n",
    "\n",
    "      res = [0] * len(nums)\n",
    "      for k, v in indices.items():\n",
    "        lv = len(v)\n",
    "        if lv == 1:\n",
    "          continue\n",
    "          \n",
    "        diffsum = diff(v)\n",
    "        for i, j in zip(v, diffsum):\n",
    "          res[i] = j\n",
    "\n",
    "      return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        dct, lst = defaultdict(list), [0] * len(nums)\n",
    "        for i, x in enumerate(nums): dct[x].append(i)\n",
    "        for k, l in dct.items(): \n",
    "            n, lpre = len(l), list(accumulate(l, initial = 0))\n",
    "            for i, x in enumerate(l): lst[x] = lpre[-1] - (lpre[i] << 1) - (n - (i << 1)) * x \n",
    "        return lst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = sum(x - a[0] for x in a)  # a[0] 到其它下标的距离之和\n",
    "            ans[a[0]] = s\n",
    "            for i in range(1, n):\n",
    "                # 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少\n",
    "                s += (i * 2 - n) * (a[i] - a[i - 1])\n",
    "                ans[a[i]] = s\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        pos = defaultdict(list) \n",
    "        \n",
    "        for i, x in enumerate(nums):\n",
    "            pos[x].append(i) \n",
    "        n = len(nums)\n",
    "        arr = [0] * n \n",
    "        def lower_bound(x, vec):\n",
    "            l, r = 0, len(vec)\n",
    "            while l < r:\n",
    "                mid = (l +r ) >> 1\n",
    "                if vec[mid] >= x:\n",
    "                    r = mid \n",
    "                else:\n",
    "                    l = mid + 1\n",
    "            return l \n",
    "            \n",
    "        for x, vec in pos.items():\n",
    "            m = len(vec)\n",
    "            f = [0] * (m + 1)\n",
    "            for i in range(1, m+1):\n",
    "                f[i] = f[i-1] + vec[i-1] \n",
    "            for i in range(m):\n",
    "                idx = lower_bound(vec[i], vec)\n",
    "                arr[vec[i]] = f[m] - vec[i] * m + 2 * (vec[i] * idx - f[idx])\n",
    "        \n",
    "        return arr "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        pos = defaultdict(list) \n",
    "        \n",
    "        for i, x in enumerate(nums):\n",
    "            pos[x].append(i) \n",
    "        n = len(nums)\n",
    "        arr = [0] * n \n",
    "        for x, vec in pos.items():\n",
    "            m = len(vec)\n",
    "            f = [0] * (m + 1)\n",
    "            for i in range(1, m+1):\n",
    "                f[i] = f[i-1] + vec[i-1] \n",
    "            for i in range(m):\n",
    "                arr[vec[i]] = f[m] - vec[i] * m + 2 * (vec[i] * i - f[i])\n",
    "        \n",
    "        return arr "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = sum(x-a[0] for x in a)\n",
    "            ans[a[0]] = s\n",
    "            for i in range(1,n):\n",
    "                s += (i * 2 - n) * (a[i] - a[i - 1])\n",
    "                ans[a[i]] = s\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        #n=len(nums)\n",
    "        #ans=[0]*n\n",
    "        #group=defaultdict(list)\n",
    "        ##{1:[0,2,3];3:[1];2:[4]}\n",
    "        #for i,x in enumerate(nums):\n",
    "        #    group[x].append(i)\n",
    "        #for a in group.values():\n",
    "        #    s=list(accumulate(a,initial=0))\n",
    "        #    #[0,0,2,5]\n",
    "        #    m=len(a)\n",
    "        #    for i in a:\n",
    "        #        b=bisect_left(a,i)\n",
    "        #        left=i*b-s[b]\n",
    "        #        right=s[m]-s[b]-i*(m-b)\n",
    "        #        ans[i]=left+right\n",
    "        #return ans\n",
    "\n",
    "        n=len(nums)\n",
    "        ans=[0]*n\n",
    "        group=defaultdict(list)\n",
    "        #{1:[0,2,3];3:[1];2:[4]}\n",
    "        for i,x in enumerate(nums):\n",
    "            group[x].append(i)\n",
    "        for a in group.values():\n",
    "            m=len(a)\n",
    "            s=0\n",
    "            first=a[0]\n",
    "            for i in range(1,m):\n",
    "                s += abs(first-a[i])\n",
    "            ans[a[0]]=s\n",
    "            for i in range(1,m):\n",
    "                ans[a[i]] = s+(2*i-m)*(a[i]-a[i-1])\n",
    "                s += (2*i-m)*(a[i]-a[i-1])\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        groups= defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            groups[num].append(i)\n",
    "        \n",
    "        ans = [0] * n\n",
    "        for idxs in groups.values():\n",
    "            m = len(idxs)\n",
    "            s = sum(j - idxs[0] for j in idxs)\n",
    "            ans[idxs[0]] = s\n",
    "            for i in range(1, m):\n",
    "                s -= (idxs[i] - idxs[i - 1]) * (m - 2 * i)\n",
    "                ans[idxs[i]] = s\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n=len(nums)\n",
    "        groups=defaultdict(list)\n",
    "        for i,x in enumerate(nums):\n",
    "            groups[x].append(i)\n",
    "        arr=[0]*n\n",
    "        for group in groups.values():\n",
    "            n=len(group)\n",
    "            g0=sum(i-group[0] for i in group)\n",
    "            arr[group[0]]=g0\n",
    "            for i in range(1,n):\n",
    "                g0+=(2*i-n)*(group[i]-group[i-1])\n",
    "                arr[group[i]]=g0\n",
    "        return arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = sum(x - a[0] for x in a)  # a[0] 到其它下标的距离之和\n",
    "            ans[a[0]] = s\n",
    "            for i in range(1, n):\n",
    "                # 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少\n",
    "                s += (i * 2 - n) * (a[i] - a[i - 1])\n",
    "                ans[a[i]] = s\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List, Dict\n",
    "from collections import defaultdict\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        d: Dict[int, list] = defaultdict(list)\n",
    "        for i in range(len(nums)):\n",
    "            d[nums[i]].append(i)\n",
    "        for k, v in d.items():\n",
    "            if len(v) == 1:\n",
    "                nums[v[0]] = 0\n",
    "                continue\n",
    "            pre_sum = 0\n",
    "            aft_sum = sum(v)\n",
    "            for i in range(len(v)):\n",
    "                aft_sum -= v[i]\n",
    "                nums[v[i]] = abs(pre_sum - v[i] * i) + abs(aft_sum - v[i] * (len(v) - i - 1))\n",
    "                pre_sum += v[i]\n",
    "        return nums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        dct, lst = defaultdict(list), [0] * len(nums)\n",
    "        for i, x in enumerate(nums): dct[x].append(i)\n",
    "        for k, l in dct.items(): \n",
    "            n, lpre = len(l), list(accumulate(l, initial = 0))\n",
    "            for i, x in enumerate(l): lst[x] = lpre[-1] - (lpre[i] << 1) - (n - (i << 1)) * x \n",
    "        return lst\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        d = defaultdict(list)\n",
    "        for i,x in enumerate(nums):\n",
    "            d[x].append(i)\n",
    "\n",
    "        ans =[0]*len(nums) \n",
    "        for a in d.values():\n",
    "            n =len(a)\n",
    "            s = sum(x-a[0] for x in a)\n",
    "            ans[a[0]]=s\n",
    "\n",
    "            for i in range(1,n):\n",
    "                s += (2*i-n)*(a[i]-a[i-1])\n",
    "                ans[a[i]]=s\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    # # 相同元素分组+考虑增量\n",
    "    # https://leetcode.cn/problems/sum-of-distances/solutions/2216380/zhao-ban-2602-ti-by-endlesscheng-6pbi/\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = sum(x - a[0] for x in a)  # a[0] 到其它下标的距离之和\n",
    "            ans[a[0]] = s\n",
    "            for i in range(1, n):\n",
    "                # 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少\n",
    "                s += (i * 2 - n) * (a[i] - a[i - 1])\n",
    "                ans[a[i]] = s\n",
    "                \n",
    "        return ans\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "\n",
    "        groups = defaultdict(list)\n",
    "\n",
    "        for i, x in enumerate(nums):\n",
    "\n",
    "            groups[x].append(i)  \n",
    "\n",
    "        ans = [0] * len(nums)\n",
    "\n",
    "        for a in groups.values():\n",
    "\n",
    "            n = len(a)\n",
    "\n",
    "            s = sum(x - a[0] for x in a)\n",
    "\n",
    "            ans[a[0]] = s\n",
    "\n",
    "            for i in range(1, n):\n",
    "\n",
    "\n",
    "                s += (i * 2 - n) * (a[i] - a[i - 1])\n",
    "\n",
    "                ans[a[i]] = s\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = sum(x - a[0] for x in a)  # a[0] 到其它下标的距离之和\n",
    "            ans[a[0]] = s\n",
    "            for i in range(1, n):\n",
    "                # 从计算 a[i-1] 到计算 a[i]，考虑 s 增加了多少\n",
    "                s += (i * 2 - n) * (a[i] - a[i - 1])\n",
    "                ans[a[i]] = s\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        idxs = defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            idxs[num].append(i)\n",
    "        \n",
    "        n = len(nums)\n",
    "        ans = [0] * n\n",
    "        for arr in idxs.values():\n",
    "            l = len(arr)\n",
    "            s = sum(idx-arr[0] for idx in arr)\n",
    "            ans[arr[0]] = s \n",
    "            for i in range(1, l):\n",
    "                s += (2*i-l) * (arr[i]-arr[i-1])\n",
    "                ans[arr[i]] = s\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        d_nums = defaultdict(list)\n",
    "        \n",
    "        n = len(nums)\n",
    "        for i in range(n):\n",
    "            d_nums[nums[i]].append(i)\n",
    "            #ans[nums[i]].append(0)\n",
    "        ans = [0] * n \n",
    "        for k, vec in d_nums.items():\n",
    "            m = len(vec)\n",
    "            f = [0] * (m)\n",
    "            f[0] = vec[0]\n",
    "            for i in range(1, m):\n",
    "                f[i] = f[i-1] + vec[i]\n",
    "\n",
    "            a = 0\n",
    "            for i in range(m):\n",
    "                a = (f[m-1] - f[i] - (m-1 - i) * vec[i]) + ((i+1) * vec[i] - f[i])\n",
    "                ans[vec[i]] = a \n",
    "\n",
    "        return ans \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        pos = defaultdict(list)\n",
    "        n = len(nums)\n",
    "        ret = [0] * n\n",
    "        for i in range(n):\n",
    "            x = nums[i]\n",
    "            pos[x].append(i)\n",
    "        for v in pos.values():\n",
    "            s = 0\n",
    "            x = v[0]\n",
    "            for i in range(1, len(v)):\n",
    "                s += v[i] - x\n",
    "            ret[v[0]] = s\n",
    "            for i in range(1, len(v)):\n",
    "                d = v[i] - v[i - 1]\n",
    "                s +=  (2 * i - len(v)) * d \n",
    "                ret[v[i]] = s\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = list(accumulate(a, initial=0))\n",
    "            for j, p in enumerate(a):\n",
    "                left = p * j - s[j]\n",
    "                right = s[n] - s[j] - p * (n - j)\n",
    "                ans[p] = left + right\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        pos = defaultdict(list)\n",
    "        for i in range(len(nums)):\n",
    "            pos[nums[i]].append(i)\n",
    "        ans = [0] * len(nums)\n",
    "        for a in pos.values():\n",
    "            n = len(a)\n",
    "            s = list(accumulate(a, initial=0))  # 前缀和\n",
    "            for j, target in enumerate(a):\n",
    "                left = target * j - s[j]\n",
    "                right = s[n] - s[j+1] - target * (n - j - 1)  # 绿色面积\n",
    "                ans[target] = left + right\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from itertools import *\n",
    "from collections import *\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups=defaultdict(list)\n",
    "        for i,num in enumerate(nums):\n",
    "            groups[num].append(i)\n",
    "        arr=[0]*len(nums)\n",
    "        for _,group in groups.items():\n",
    "            pre_sum=list(accumulate(group,initial=0))\n",
    "            for i, num in enumerate(group):\n",
    "                left=i*num-pre_sum[i]\n",
    "                right=pre_sum[-1]-pre_sum[i]-num*(len(group)-i)\n",
    "                arr[num]=left+right\n",
    "        return arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        d = defaultdict(list)\n",
    "        for i,x in enumerate(nums):\n",
    "            d[x].append(i)\n",
    "\n",
    "        ans =[0]*len(nums) \n",
    "        for a in d.values():\n",
    "            n =len(a)\n",
    "\n",
    "            pre = list(accumulate(a,initial=0))\n",
    "\n",
    "            for i, x in enumerate(a):\n",
    "                left = x*i -pre[i]\n",
    "                right = pre[n]-pre[i] -x *(n-i)\n",
    "                ans[x] = left +right\n",
    "\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        d = defaultdict(list)\n",
    "        for i in range(len(nums)):\n",
    "            d[nums[i]].append(i)\n",
    "        res = [0] * len(nums)\n",
    "        for k, arr in d.items():\n",
    "            accu = [0]\n",
    "            for i, x in enumerate(arr):\n",
    "                accu.append(accu[i] + x)\n",
    "            for i, x in enumerate(arr):\n",
    "                v1 = x * i - accu[i]\n",
    "                v2 = accu[-1] - accu[i] - x * (len(arr) - i)\n",
    "                res[x] = v1 + v2\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# class Solution:\n",
    "#     def distance(self, nums: List[int]) -> List[int]:\n",
    "#         cnt = defaultdict(list)\n",
    "#         for i,num in enumerate(nums):\n",
    "#             cnt[num].append(i)\n",
    "        \n",
    "#         ans = []\n",
    "#         for i in range(len(nums)):\n",
    "#             tmp = cnt[nums[i]]\n",
    "#             total = 0\n",
    "#             for t in tmp:\n",
    "#                 total += abs(i - t)\n",
    "#             ans.append(total)\n",
    "#         return ans\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = list(accumulate(a, initial=0))  # 前缀和\n",
    "            for j, target in enumerate(a):\n",
    "                left = target * j - s[j]  # 蓝色面积\n",
    "                right = s[n] - s[j] - target * (n - j)  # 绿色面积\n",
    "                ans[target] = left + right\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        dct, lst = defaultdict(list), [0] * len(nums)\n",
    "        for i, x in enumerate(nums): dct[x].append(i)\n",
    "        for k, l in dct.items(): \n",
    "            n, lpre = len(l), list(accumulate(l, initial = 0))\n",
    "            for i, x in enumerate(l):\n",
    "                blue = i*x-lpre[i]\n",
    "                grenn = lpre[-1]-lpre[i]-x*(n-i)\n",
    "                lst[x]=blue+grenn\n",
    "        return lst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = list(accumulate(a, initial = 0))\n",
    "            for j ,target in enumerate(a):\n",
    "                left = target * j - s[j]\n",
    "                right = s[n] - s[j] - target * (n - j)\n",
    "                ans[target] = left + right\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        g = defaultdict(list)\n",
    "        for i, v in enumerate(nums):\n",
    "            g[v].append(i)\n",
    "        ans = [0] * len(nums)\n",
    "        for arr in g.values():\n",
    "            m = len(arr)\n",
    "            s = [0] * (m + 1)\n",
    "            for i in range(m):\n",
    "                s[i + 1] = s[i] + arr[i]\n",
    "            for i, v in enumerate(arr):\n",
    "                l = i * v - s[i]\n",
    "                r = s[m] - s[i] - v * (m - i)\n",
    "                ans[v] = l + r\n",
    "        return ans\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        d = defaultdict(list)\n",
    "        for i in range(len(nums)):\n",
    "            d[nums[i]].append(i)\n",
    "        res = [0] * len(nums)\n",
    "        for k, arr in d.items():\n",
    "            accu = [0]\n",
    "            for i, x in enumerate(arr):\n",
    "                accu.append(accu[i] + x)\n",
    "            for i, x in enumerate(arr):\n",
    "                v1 = x * i - accu[i]\n",
    "                v2 = accu[-1] - accu[i] - x * (len(arr) - i)\n",
    "                res[x] = v1 + v2\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def distance(self, nums: list[int]) -> list[int]:\r\n",
    "        # from collections import defaultdict\r\n",
    "        dict_list = defaultdict(list)\r\n",
    "        for i, num in enumerate(nums):\r\n",
    "            dict_list[num].append(i)\r\n",
    "        # print(dict_list)\r\n",
    "        res = [0] * len(nums)\r\n",
    "        for d in dict_list.values():\r\n",
    "            n = len(d)\r\n",
    "            if n == 1:\r\n",
    "                continue\r\n",
    "            pre_sum = [0] * n\r\n",
    "            pre_sum[0] = d[0]\r\n",
    "            for i in range(1, n):\r\n",
    "                pre_sum[i] = pre_sum[i - 1] + d[i]\r\n",
    "            # print(pre_sum)\r\n",
    "            for i in range(n):\r\n",
    "                if i == 0:\r\n",
    "                    res[d[i]] += pre_sum[-1] - pre_sum[i] - (n - 1 - i) * d[i]\r\n",
    "                elif i == n - 1:\r\n",
    "                    res[d[i]] += i * d[i] - pre_sum[i - 1]\r\n",
    "                else:\r\n",
    "                    res[d[i]] += i * d[i] - pre_sum[i - 1] + pre_sum[-1] - pre_sum[i] - (n - 1 - i) * d[i]\r\n",
    "        # print(res)\r\n",
    "        return res\r\n",
    "\r\n",
    "\r\n",
    "# s = Solution()\r\n",
    "# s.distance([1,3,1,1,2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n=len(nums)\n",
    "        ref = collections.defaultdict(list)\n",
    "        for i in range(n):\n",
    "            ref[nums[i]].append(i)\n",
    "        #print(ref)\n",
    "\n",
    "        ans = [0]*n\n",
    "        from itertools import accumulate\n",
    "        for k in ref:\n",
    "            length = len(ref[k])\n",
    "            acc = list(accumulate(ref[k],initial=0))\n",
    "\n",
    "            #print(acc)\n",
    "            for j,index in enumerate(ref[k]):\n",
    "                left = index*j - acc[j]\n",
    "                right = acc[length]-acc[j]-index*(len(ref[k])-j)\n",
    "               # print(left,right)\n",
    "                ans[index] = left+right\n",
    "        return ans\n",
    "\n",
    "        return []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "from itertools import accumulate\n",
    "from typing import List\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = list(accumulate(a, initial=0))  # 前缀和\n",
    "            for j, target in enumerate(a):\n",
    "                left = target * j - s[j]  # 蓝色面积\n",
    "                right = s[n] - s[j] - target * (n - j)  # 绿色面积\n",
    "                ans[target] = left + right\n",
    "        return ans\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        # map=defaultdict(list)\n",
    "        # for i,num in enumerate(nums):\n",
    "        #     map[num].append(i)\n",
    "        # # print(map)\n",
    "        # map1=defaultdict(list)\n",
    "        # for key in map:\n",
    "        #     n=len(map[key])\n",
    "        #     pre_sum=[0]*(n+1)\n",
    "        #     for i in range(1,n+1):\n",
    "        #         pre_sum[i]=pre_sum[i-1]+map[key][i-1]\n",
    "        #     map1[key]=pre_sum\n",
    "        # # print(map1)\n",
    "        # arr=[]\n",
    "        # for i,num in enumerate(nums):\n",
    "        #     a=len(map[num])\n",
    "        #     idx=bisect.bisect_left(map[num], i)\n",
    "            \n",
    "        #     arr.append(idx*i-(map1[num][idx]-map1[num][0])+map1[num][-1]-map1[num][idx+1]-(a-idx-1)*i)\n",
    "\n",
    "        # return arr\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = list(accumulate(a, initial=0))  # 前缀和\n",
    "            for j, target in enumerate(a):\n",
    "                left = target * j - s[j]  # 蓝色面积\n",
    "                right = s[n] - s[j] - target * (n - j)  # 绿色面积\n",
    "                ans[target] = left + right\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        c=collections.defaultdict(list)\n",
    "        for i,n in enumerate(nums):\n",
    "            c[n].append(i)\n",
    "        ans=[0]*len(nums)\n",
    "        for key in c:\n",
    "            ls=c[key]\n",
    "            n=len(ls)\n",
    "            pre=list(accumulate(ls,initial=0))\n",
    "            for i,k in enumerate(ls):\n",
    "                l=k*i-pre[i]\n",
    "                r=pre[n]-pre[i]-k*(n-i)\n",
    "                ans[k]=l+r\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        groups = defaultdict(list)\n",
    "        for i, x in enumerate(nums):\n",
    "            groups[x].append(i)  # 相同元素分到同一组，记录下标\n",
    "        ans = [0] * len(nums)\n",
    "        for a in groups.values():\n",
    "            n = len(a)\n",
    "            s = list(accumulate(a, initial=0))  # 前缀和\n",
    "            for j, target in enumerate(a):\n",
    "                left = target * j - s[j]  # 蓝色面积\n",
    "                right = s[n] - s[j] - target * (n - j)  # 绿色面积\n",
    "                ans[target] = left + right\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        ans = [0] * len(nums)\n",
    "        g = defaultdict(list)\n",
    "        for i,x in enumerate(nums):\n",
    "            g[x].append(i)\n",
    "\n",
    "        for n in g.values():\n",
    "            m = len(n)\n",
    "            s = list(accumulate(n,initial = 0))\n",
    "            for j,target in enumerate(n):\n",
    "                l = target * j - s[j]\n",
    "                r = s[m] - s[j] - target * (m-j)\n",
    "                ans[target] = l+r\n",
    "        \n",
    "\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n=len(nums)\n",
    "        ans=[0]*n\n",
    "        group=defaultdict(list)\n",
    "        #{1:[0,2,3];3:[1];2:[4]}\n",
    "        for i,x in enumerate(nums):\n",
    "            group[x].append(i)\n",
    "        for a in group.values():\n",
    "            s=list(accumulate(a,initial=0))\n",
    "            #[0,0,2,5]\n",
    "            m=len(a)\n",
    "            for i in a:\n",
    "                b=bisect_left(a,i)\n",
    "                left=i*b-s[b]\n",
    "                right=s[m]-s[b]-i*(m-b)\n",
    "                ans[i]=left+right\n",
    "        return ans\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        mp=defaultdict(list)\n",
    "        arr=[0]*len(nums)\n",
    "        for i in range(len(nums)):\n",
    "            mp[nums[i]].append(i)\n",
    "        for key in mp.keys():\n",
    "            n=len(mp[key])\n",
    "            suf=[0]*(n+1)\n",
    "            for i in range(n):\n",
    "                suf[i+1]=suf[i]+mp[key][i]\n",
    "            for i in range(n):\n",
    "                left=mp[key][i]*i-suf[i]\n",
    "                right=suf[n]-suf[i]-mp[key][i]*(n-i)\n",
    "                arr[mp[key][i]]=left+right\n",
    "        return arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        mp=defaultdict(list)\n",
    "        arr=[0]*len(nums)\n",
    "        for i in range(len(nums)):\n",
    "            mp[nums[i]].append(i)\n",
    "        for key in mp.keys():\n",
    "            n=len(mp[key])\n",
    "            suf=[0]*(n+1)\n",
    "            for i in range(n):\n",
    "                suf[i+1]=suf[i]+mp[key][i]\n",
    "            for i in range(n):\n",
    "                left=mp[key][i]*i-suf[i]\n",
    "                right=suf[n]-suf[i+1]-mp[key][i]*(n-i-1)\n",
    "                arr[mp[key][i]]=left+right\n",
    "        return arr\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        res = [0] * n\n",
    "        mp = defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            mp[num].append(i) \n",
    "        \n",
    "        for num, idxs in mp.items():\n",
    "            lenIdxs = len(idxs)\n",
    "            leftsum = []\n",
    "            p = 0\n",
    "            for idx in idxs:\n",
    "                p += idx \n",
    "                leftsum.append(p)\n",
    "            rightsum = [0] * len(idxs)\n",
    "            p = 0 \n",
    "            for i in range(len(idxs)-1, -1, -1):\n",
    "                p += idxs[i]\n",
    "                rightsum[i] = p \n",
    "            print(num, leftsum, rightsum)\n",
    "            for i, curidx in enumerate(idxs):\n",
    "                curres = 0\n",
    "                if i - 1 >= 0:\n",
    "                    # curidx * (num of idx to the left) - leftsum\n",
    "                    curres += curidx * i - leftsum[i-1]\n",
    "                if i + 1 < lenIdxs:\n",
    "                     # rightsum  - curidx * (num of idx to the right) \n",
    "                    curres += rightsum[i+1] - curidx * (lenIdxs - i - 1) \n",
    "    \n",
    "                res[curidx] = curres\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        ans = [0] * n\n",
    "        d = collections.defaultdict(list)\n",
    "        for i, nu in enumerate(nums):\n",
    "            d[nu].append(i)\n",
    "        for k, vl in d.items():\n",
    "            n_vl = len(vl)\n",
    "            if n_vl == 1:\n",
    "                continue\n",
    "            s = sum(vl)\n",
    "            first = s - n_vl * vl[0]\n",
    "            ans[vl[0]] = first\n",
    "            # print(vl)\n",
    "            for i in range(1, n_vl):\n",
    "                first -= (vl[i] - vl[i - 1]) * (n_vl - i)\n",
    "                first += i * (vl[i] - vl[i - 1])\n",
    "                # print(i, first)\n",
    "                ans[vl[i]] = first\n",
    "        return ans\n",
    "                \n",
    "            \n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        res = [0] * n\n",
    "        mp = defaultdict(list)\n",
    "        for i, num in enumerate(nums):\n",
    "            mp[num].append(i) \n",
    "        for num, idxs in mp.items():\n",
    "            lenIdxs = len(idxs)\n",
    "            leftsum = []\n",
    "            p = 0\n",
    "            for idx in idxs:\n",
    "                p += idx \n",
    "                leftsum.append(p)\n",
    "            rightsum = [0] * len(idxs)\n",
    "            p = 0 \n",
    "            for i in range(len(idxs)-1, -1, -1):\n",
    "                p += idxs[i]\n",
    "                rightsum[i] = p \n",
    "            for i, curidx in enumerate(idxs):\n",
    "                curres = 0\n",
    "                if i - 1 >= 0:\n",
    "                    # curidx * (num of idx to the left) - leftsum\n",
    "                    curres += curidx * i - leftsum[i-1]\n",
    "                if i + 1 < lenIdxs:\n",
    "                     # rightsum  - curidx * (num of idx to the right) \n",
    "                    curres += rightsum[i+1] - curidx * (lenIdxs - i - 1) \n",
    "    \n",
    "                res[curidx] = curres\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        def func(nums):\n",
    "            memo = dict()\n",
    "            ans = []\n",
    "            for idx,key in enumerate(nums):\n",
    "                if key not in memo:\n",
    "                    memo[key] = [idx,1,0]\n",
    "                else:\n",
    "                    memo[key][2] += memo[key][1]*(idx - memo[key][0])\n",
    "                    memo[key][1] += 1\n",
    "                    memo[key][0] = idx\n",
    "                ans.append(memo[key][2])\n",
    "            return ans\n",
    "        res1 = func(nums)\n",
    "        res2 = func(nums[::-1])[::-1]\n",
    "        return [x + y for x, y in zip(res1, res2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        def func(nums):\n",
    "            memo = dict()\n",
    "            ans = []\n",
    "            for idx,key in enumerate(nums):\n",
    "                if key not in memo:\n",
    "                    memo[key] = [idx,1,0]\n",
    "                else:\n",
    "                    memo[key][2] += memo[key][1]*(idx - memo[key][0])\n",
    "                    memo[key][1] += 1\n",
    "                    memo[key][0] = idx\n",
    "                ans.append(memo[key][2])\n",
    "            return ans\n",
    "        res1 = func(nums)\n",
    "        res2 = func(nums[::-1])[::-1]\n",
    "        ans = []\n",
    "        for i in range(len(nums)):\n",
    "            ans.append(res1[i] + res2[i])\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        s=len(nums)\n",
    "        arr=[]\n",
    "        tot={}\n",
    "        pp={}\n",
    "        for i in range(s):\n",
    "            if nums[i] not in tot:\n",
    "                tot[nums[i]]=[i]\n",
    "                pp[nums[i]]=i\n",
    "            else:\n",
    "                tot[nums[i]].append(i)\n",
    "                pp[nums[i]]+=i\n",
    "        p=list(tot.values())\n",
    "        q=list(pp.values())\n",
    "        for i in range(len(p)):\n",
    "            t=len(p[i])\n",
    "            ans=q[i]-t*p[i][0]\n",
    "            nums[p[i][0]]=ans\n",
    "            b=q[i]-p[i][0]\n",
    "            a=p[i][0]\n",
    "            for j in range(1,t):\n",
    "                ans=b-(t-j*2)*p[i][j]-a\n",
    "                nums[p[i][j]]=ans\n",
    "                a+=p[i][j]\n",
    "                b-=p[i][j]\n",
    "        return nums\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        s=len(nums)\n",
    "        arr=[]\n",
    "        tot={}\n",
    "        pp={}\n",
    "        for i in range(s):\n",
    "            if nums[i] not in tot:\n",
    "                tot[nums[i]]=[i]\n",
    "                pp[nums[i]]=i\n",
    "            else:\n",
    "                tot[nums[i]].append(i)\n",
    "                pp[nums[i]]+=i\n",
    "        p=list(tot.values())\n",
    "        q=list(pp.values())\n",
    "        for i in range(len(p)):\n",
    "            t=len(p[i])\n",
    "            ans=q[i]-t*p[i][0]\n",
    "            nums[p[i][0]]=ans\n",
    "            b=q[i]-p[i][0]\n",
    "            a=p[i][0]\n",
    "            for j in range(1,t):\n",
    "                ans=b-(t-j*2)*p[i][j]-a\n",
    "                nums[p[i][j]]=ans\n",
    "                a+=p[i][j]\n",
    "                b-=p[i][j]\n",
    "        return nums\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n=len(nums)\n",
    "\n",
    "        cnt=defaultdict(list)\n",
    "        tot=defaultdict()\n",
    "        for i,e in enumerate(nums):\n",
    "            cnt[e].append(i)\n",
    "            cur=i-cnt[e][0]\n",
    "            tot[e] =cur+(0 if cur==0 else tot[e])\n",
    "        res=[0]*n\n",
    "        for k,v in cnt.items():\n",
    "            res[v[0]]=tot[k]\n",
    "            pre=tot[k]\n",
    "            for i in range(1,len(v)):\n",
    "                diff=v[i]-v[i-1]\n",
    "                x=i+i-len(v)\n",
    "                res[v[i]]=pre+x*diff\n",
    "                pre=res[v[i]]\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        n = len(nums)\n",
    "        pres = Counter()\n",
    "        pcnt = Counter()\n",
    "        suff = Counter()\n",
    "        cnt = Counter()\n",
    "        for i in range(n - 1, -1, -1):\n",
    "            cnt[nums[i]] += 1\n",
    "            suff[nums[i]] += i\n",
    "        ans = [0] * n\n",
    "        for i, x in enumerate(nums):\n",
    "            suff[x] -= i\n",
    "            pcnt[x] += 1\n",
    "            \n",
    "            if cnt[x] > 1:\n",
    "                ans[i] = suff[x] - (cnt[x] - pcnt[x]) * i + (pcnt[x] - 1) * i - pres[x]\n",
    "            pres[x] += i\n",
    "\n",
    "        return ans "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "\n",
    "        mp = collections.defaultdict(int)\n",
    "        cnts = collections.Counter(nums)\n",
    "\n",
    "        for i, n in enumerate(nums):\n",
    "            mp[n] += i \n",
    "\n",
    "        prefix_sum = collections.defaultdict(int)\n",
    "        prefix_cnt = collections.Counter()\n",
    "        res = [0] * len(nums)\n",
    "        for i, n in enumerate(nums):\n",
    "            pre_sum = prefix_sum.get(n, 0)\n",
    "            pre_cnt = prefix_cnt.get(n, 0)\n",
    "            suf_sum = mp.get(n, 0) - pre_sum - i \n",
    "            suf_cnt = cnts.get(n, 0) - pre_cnt - 1\n",
    "\n",
    "            res[i] = (i * pre_cnt - pre_sum) + (suf_sum - i * suf_cnt)\n",
    "\n",
    "            prefix_sum[n] += i \n",
    "            prefix_cnt[n] += 1\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distance(self, nums: List[int]) -> List[int]:\n",
    "        map=defaultdict(list)\n",
    "        for i,num in enumerate(nums):\n",
    "            map[num].append(i)\n",
    "        # print(map)\n",
    "        map1=defaultdict(list)\n",
    "        for key in map:\n",
    "            n=len(map[key])\n",
    "            pre_sum=[0]*(n+1)\n",
    "            for i in range(1,n+1):\n",
    "                pre_sum[i]=pre_sum[i-1]+map[key][i-1]\n",
    "            map1[key]=pre_sum\n",
    "        # print(map1)\n",
    "        arr=[]\n",
    "        for i,num in enumerate(nums):\n",
    "            a=len(map[num])\n",
    "            idx=bisect.bisect_left(map[num], i)\n",
    "            \n",
    "            arr.append(idx*i-(map1[num][idx]-map1[num][0])+map1[num][-1]-map1[num][idx+1]-(a-idx-1)*i)\n",
    "\n",
    "        return arr"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
